from PIL import Image
import os

origin_path = 'codes/'
new_path = 'parsed/'
print(origin_path)

for image in os.listdir(origin_path)[:5]:
    im = Image.open(origin_path + image)
    width, height = im.size

    # 获取图片中的颜色, 返回列表 [(counts, color)...]
    color_info = im.getcolors(width*height)
    # 按照计数从大到小排列颜色, 那么颜色计数最多的是背景
    # 接下来排名第 2 到其他的是字符了或其他干扰项了
    sort_color = sorted(color_info, key = lambda x: x[0], reverse = True)
    # 打印一下, 发现只有三项, 那么下面的循环就 range(1, 3)
    #  print(sort_color)

    # 根据颜色, 提取同颜色的字符, 重新放到新建的白色背景 image 对象上.
    char_dict = {}
    for i in range(1, 3):
        im2 = Image.new('RGB', im.size, (255, 255, 255))
        for x in range(im.size[0]):
            for y in range(im.size[1]):
                if im.getpixel((x, y)) == sort_color[i][1]:
                    im2.putpixel((x, y), (0, 0, 0))
                else:
                    im2.putpixel((x, y), (255, 255, 255))
        im2.save(new_path + str(i) + '-' + image)
    print('成功处理图片{}'.format(image))

